"""
DeepCrack模型训练器 - 用于道路裂缝检测的深度学习模型训练框架

核心功能：
1. 封装完整的训练流程（前向传播、损失计算、反向传播、优化器更新）
2. 实现多尺度预测结果的融合与评估
3. 提供验证集测试和精度计算功能
4. 集成可视化工具和模型保存功能

典型应用场景：
- 基于DeepCrack网络的裂缝检测模型训练
- 多尺度特征融合的语义分割任务
- 二分类像素级预测任务

技术特点：
- 使用带权重调整的BCEWithLogitsLoss处理类别不平衡
- 支持Adam/SGD优化器灵活切换
- 实时记录训练损失和精度指标
- 集成Visdom可视化监控
"""

import torch
from tools.checkpointer import Checkpointer
from tools.visdomTools import Visualizer
from torch import nn

from config.config import Config as cfg


def get_optimizer(model):
    """
    优化器工厂函数
    作用:根据损失更新参数
    参数：
        model (nn.Module): 待优化的模型

    返回：
        torch.optim.Optimizer: 根据配置选择的优化器实例

    实现细节：
        - 通过cfg.use_adam配置切换Adam/SGD优化器
        - 权重衰减(weight_decay)用于L2正则化
        - SGD优化器使用动量(momentum)加速收敛
    """
    if cfg.use_adam:
        # Adam:  Adam优化器 考虑参数的历史信息 更加智能的调整学习率
        return torch.optim.Adam(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
    else:
        # SGD:  SGD优化器 沿着损失函数下降最快的方向进行调整参数
        return torch.optim.SGD(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay, momentum=cfg.momentum)


class DeepCrackTrainer(nn.Module):
    """
    DeepCrack模型训练器

    初始化参数：
        model (nn.Module): 要训练的DeepCrack模型实例

    主要组件：
        - Visualizer: 训练过程可视化工具
        - Checkpointer: 模型保存与加载工具
        - BCEWithLogitsLoss: 带sigmoid的二元交叉熵损失
    """

    def __init__(self, model,is_visible=True):
        """
        :param model:
        :param is_visible:  是否可视化
        """
        # 继承父类初始化
        super(DeepCrackTrainer, self).__init__()
        # 可视化工具初始化（用于监控训练过程） 如果isVisible为True 就初始化
        if is_visible:
            self.vis = Visualizer(env=cfg.vis_env)
        else:
            self.vis = None
        # 待训练模型
        self.model = model
        # 模型保存工具（支持最大保存数量限制）
        self.saver = Checkpointer(cfg.name, cfg.saver_path,
                                  overwrite=False, verbose=True,
                                  timestamp=True, max_queue=cfg.max_save)
        # 优化器（Adam/SGD）   作用:根据损失调整参数
        self.optimizer = get_optimizer(self.model)
        # 迭代计数器
        self.iter_counter = 0

        # -------------------- 损失函数配置 ---------------------
        # 带权重调整的二元交叉熵损失（处理正负样本不平衡）
        self.mask_loss = nn.BCEWithLogitsLoss(
            # reduction: 损失计算方式 mean 代表 均值
            reduction='mean',
            # pos_weight 是一个大于 0 的数，专门用来放大正样本的损失
            pos_weight=torch.tensor([cfg.pos_pixel_weight], dtype=torch.float, device='cuda')
        )

        # 日志记录字典
        self.log_loss = {}  # 记录各尺度损失值
        self.log_acc = {}  # 记录精度指标

    def train_op(self, input, target):
        """
        单次训练操作

        参数：
            input (Tensor): 输入图像 [B,C,H,W]
                                    [批次大小、通道数、高度、宽度]
            target (Tensor): 目标掩码 [B,1,H,W]
                                    [批次大小、1、高度、宽度]

        返回：
            Tuple[Tensor]: 各尺度预测结果（6个输出）

        处理流程：
            1. 梯度清零
            2. 前向传播获取多尺度预测
            3. 计算各尺度损失（加权求和）
            4. 反向传播与参数更新
            5. 记录损失指标
        """
        # 梯度清零 原因：梯度计算会累加 新的一轮要清零
        self.optimizer.zero_grad()

        # 前向传播（获取主输出和5个辅助输出）
        """
            pred_output：最终预测结果（如裂缝掩码）
            pred_fuse1 到 pred_fuse5：不同层级的特征融合输出（从浅层到深层）
        最终输出：
    output  是经过一系列处理后的最终预测结果。它是通过对各级特征进行融合得到的。
            具体过程是先对编码器和解码器的各级特征进行条带池化，然后使用 Fuse 模块进行融合，
            最后通过 Conv3X3 模块输出预测结果。这个输出的形状和维度取决于输入数据的形状、网络结构以及模型的设计。
            在 DeepCrack_uu_Sp_3 类的 forward 方法中，最终输出 output 是对输入图像的预测结果，
            用于判断图像中是否存在裂缝等目标，
            其通道数为 1（由 self.final = Conv3X3(5, 1) 决定），代表二分类（有裂缝或无裂缝）的预测。
        各级特征：
    fuse5： 是第 5 级的融合特征。它是将下采样得到的 down5 特征和上采样得到的 up5 特征，经过条带池化和 Fuse 模块处理后的结果。
            fuse5 包含了对图像中较大尺度信息的处理结果，有助于捕捉图像中的全局特征。
    fuse4： 第 4 级的融合特征。同样是对下采样的 down4 特征和上采样的 up4 特征进行条带池化和 Fuse 模块处理。
            相比 fuse5，fuse4 可能捕捉到更细粒度的特征，在模型中起到了连接不同尺度特征的作用。
    fuse3： 第 3 级的融合特征。处理方式与上述类似，它进一步融合了不同尺度的特征，使得模型能够学习到图像中不同层次的信息。
    fuse2： 第 2 级的融合特征。它对下采样的 down2 特征和上采样的 up2 特征进行处理，保留了图像中较细粒度的信息。
    fuse1： 第 1 级的融合特征。作为最细粒度的特征融合结果，fuse1 包含了图像中最细节的信息，与其他各级特征一起为最终的预测提供支持。
        """
        pred_output, pred_fuse5, pred_fuse4, pred_fuse3, pred_fuse2, pred_fuse1 = self.model(input)

        # 各尺度损失计算（归一化到batch size）
        """
        pred_output.view(-1, 1)
        pred_output:预测输出、view(-1, 1)代表将张量展开为1维，方便计算像素级损失
        target.view(-1, 1)
        target:目标(标签)输出
        """
        output_loss = self.mask_loss(pred_output.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse5_loss = self.mask_loss(pred_fuse5.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse4_loss = self.mask_loss(pred_fuse4.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse3_loss = self.mask_loss(pred_fuse3.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse2_loss = self.mask_loss(pred_fuse2.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse1_loss = self.mask_loss(pred_fuse1.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size

        # 多尺度损失融合（简单求和）
        total_loss = output_loss + fuse5_loss + fuse4_loss + fuse3_loss + fuse2_loss + fuse1_loss
        total_loss.backward()
        self.optimizer.step()

        self.iter_counter += 1

        # 记录损失指标（用于可视化）
        self.log_loss = {
            'total_loss': total_loss.item(),
            'output_loss': output_loss.item(),
            'fuse5_loss': fuse5_loss.item(),
            'fuse4_loss': fuse4_loss.item(),
            'fuse3_loss': fuse3_loss.item(),
            'fuse2_loss': fuse2_loss.item(),
            'fuse1_loss': fuse1_loss.item()
        }

        return pred_output, pred_fuse5, pred_fuse4, pred_fuse3, pred_fuse2, pred_fuse1

    def val_op(self, input, target):
        """
        验证操作（不更新参数）

        参数与返回：
            同train_op

        特殊处理：
            - 不执行梯度计算和参数更新
            - 用于验证集性能评估
        """
        # 前向传播
        pred_output, pred_fuse5, pred_fuse4, pred_fuse3, pred_fuse2, pred_fuse1 = self.model(input)

        # 损失计算（不除以batch_size）
        output_loss = self.mask_loss(pred_output.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse5_loss = self.mask_loss(pred_fuse5.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse4_loss = self.mask_loss(pred_fuse4.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse3_loss = self.mask_loss(pred_fuse3.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse2_loss = self.mask_loss(pred_fuse2.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size
        fuse1_loss = self.mask_loss(pred_fuse1.view(-1, 1), target.view(-1, 1)) / cfg.train_batch_size

        total_loss = output_loss + fuse5_loss + fuse4_loss + fuse3_loss + fuse2_loss + fuse1_loss

        # 记录验证损失
        self.log_loss = {
            'total_loss': total_loss.item(),
            'output_loss': output_loss.item(),
            'fuse5_loss': fuse5_loss.item(),
            'fuse4_loss': fuse4_loss.item(),
            'fuse3_loss': fuse3_loss.item(),
            'fuse2_loss': fuse2_loss.item(),
            'fuse1_loss': fuse1_loss.item()
        }

        return pred_output, pred_fuse5, pred_fuse4, pred_fuse3, pred_fuse2, pred_fuse1

    def val_out(self, input):
        """
        纯推理模式

        参数：
            input (Tensor): 输入图像

        返回：
            Tuple[Tensor]: 各尺度预测结果

        用途：
            - 测试阶段无标签预测
            - 模型部署时的前向计算
        """
        pred_output, pred_fuse5, pred_fuse4, pred_fuse3, pred_fuse2, pred_fuse1 = self.model(input)
        return pred_output, pred_fuse5, pred_fuse4, pred_fuse3, pred_fuse2, pred_fuse1

    def acc_op(self, pred, target):
        """
        精度计算操作

        参数：
            pred (Tensor): 模型预测输出 [B,1,H,W]
            target (Tensor): 真实标签 [B,1,H,W]

        计算指标：
            - 整体准确率
            - 正样本（裂缝）准确率
            - 负样本（背景）准确率

        处理流程：
            1. 通过阈值将预测概率二值化
            2. 计算各类别像素级准确率
        """
        mask = target

        # 预测值二值化（使用配置阈值）
        pred = pred
        pred[pred > cfg.acc_sigmoid_th] = 1
        pred[pred <= cfg.acc_sigmoid_th] = 0

        pred_mask = pred[:, 0, :, :].contiguous()

        # 整体准确率（所有像素）
        mask_acc = pred_mask.eq(mask.view_as(pred_mask)).sum().item() / mask.numel()
        # 正样本准确率（裂缝像素）
        mask_pos_acc = pred_mask[mask > 0].eq(mask[mask > 0]).sum().item() / (mask[mask > 0].numel() + 1)
        # 负样本准确率（背景像素）
        mask_neg_acc = pred_mask[mask < 1].eq(mask[mask < 1]).sum().item() / (mask[mask < 1].numel() + 1)

        # 记录精度指标
        self.log_acc = {
            'mask_acc': mask_acc,
            'mask_pos_acc': mask_pos_acc,
            'mask_neg_acc': mask_neg_acc,
        }